#include <iostream>
#include "TFile.h"
#include "TTree.h"
#include "TH1D.h"
#include "HelixHough.h"
#include "HelixRange.h"
#include "HelixResolution.h"
#include "VertexFitFunc.h"
#include "NewtonMinimizerGradHessian.h"
#include <sys/time.h>
#include <Eigen/LU>
#include <Eigen/Core>
#include <math.h>
#include "sPHENIXTracker.h"
#include "SimpleTrack.h"


using namespace std;
using namespace FitNewton;
using namespace Eigen;


int main(int argc, char** argv)
{
  bool print_truth = false;
  bool print_reco = false;
  
  TFile infile(argv[1]);
  TTree* etree=0;
  infile.GetObject("events", etree);
  SimpleMCEvent* inputevent=0;
  etree->SetBranchAddress("tracks", &inputevent);
  
  int nlayers = 6;
  vector<float> radii;
  radii.assign(nlayers,0.);
  radii[0]=2.5;
  radii[1]=5.0;
  radii[2]=10.0;
  radii[3]=14.0;
  radii[4]=40.0;
  radii[5]=60.0;
  vector<float> smear_xy_layer;smear_xy_layer.assign(nlayers,0);
  vector<float> smear_z_layer;smear_z_layer.assign(nlayers,0);
  float sqrt_12 = sqrt(12.);
  smear_xy_layer[0] = (50.0e-4/sqrt_12);
  smear_z_layer[0] = (425.0e-4/sqrt_12);
  smear_xy_layer[1] = (50.0e-4/sqrt_12);
  smear_z_layer[1] = (425.0e-4/sqrt_12);
  smear_xy_layer[2] = (80.0e-4/sqrt_12);
  smear_z_layer[2] = (1000.0e-4/sqrt_12);
  smear_xy_layer[3] = (80.0e-4/sqrt_12);
  smear_z_layer[3] = (1000.0e-4/sqrt_12);
  smear_xy_layer[4] = (80.0e-4/sqrt_12);
  smear_z_layer[4] = (30000.0e-4/sqrt_12);
  smear_xy_layer[5] = (80.0e-4/sqrt_12);
  smear_z_layer[5] = (30000.0e-4/sqrt_12);
  
  float kappa_max = 0.03;
  float rho_max = pow(kappa_max,1.);
  
  
  // *********************************************************
  // setup for initial 4-hit seed-finding
  
  //phi,d,kappa,dzdl,z0
  float min_iter = 2.;
  float max_iter = 9.;
  HelixRange top_range(0., 2.*M_PI,   -0.1, 0.1,   0.0, rho_max, -0.9, 0.9,   -0.1, 0.1);
  
  vector<unsigned int> onezoom(5,0);
  vector<vector<unsigned int> > zoomprofile;
  zoomprofile.assign(9,onezoom);
  for(unsigned int i=0;i<=2;++i)
  {
    zoomprofile[i][0] = 5;
    zoomprofile[i][1] = 1;
    zoomprofile[i][2] = 2;
    zoomprofile[i][3] = 4;
    zoomprofile[i][4] = 1;
  }
  for(unsigned int i=3;i<=8;++i)
  {
    zoomprofile[i][0] = 3;
    zoomprofile[i][1] = 1;
    zoomprofile[i][2] = 2;
    zoomprofile[i][3] = 4;
    zoomprofile[i][4] = 1;
  }
  sPHENIXTracker tracker(zoomprofile, 3, top_range);
  tracker.setNLayers(4);
  unsigned int max_hits_seed = 36;
  unsigned int min_hits_seed = 4;
  tracker.setRejectGhosts(false);
  tracker.setChi2Cut(10.0);
  tracker.setPrintTimings(true);
  tracker.setVerbosity(1);
  tracker.setCutOnDca(false);
  tracker.requireLayers(4);
  tracker.setSmoothBack(false);
  
  
  // *********************************************************
  // setup for 6-hit tracking based on 4-hit seeds
  vector<vector<unsigned int> > zoomprofile_2;
  zoomprofile_2.assign(9,onezoom);
  for(unsigned int i=0;i<=2;++i)
  {
    zoomprofile_2[i][0] = 5;
    zoomprofile_2[i][1] = 1;
    zoomprofile_2[i][2] = 2;
    zoomprofile_2[i][3] = 3;
    zoomprofile_2[i][4] = 1;
  }
  for(unsigned int i=3;i<=8;++i)
  {
    zoomprofile_2[i][0] = 3;
    zoomprofile_2[i][1] = 1;
    zoomprofile_2[i][2] = 2;
    zoomprofile_2[i][3] = 3;
    zoomprofile_2[i][4] = 1;
  }
  HelixRange top_range_2(0., 2.*M_PI,   -0.2, 0.2,   0.0, rho_max,   -0.9, 0.9,   -0.3, 0.3);
  sPHENIXTracker tracker_seeded(zoomprofile_2, 3, top_range_2);
  unsigned int max_hits = 36;
  unsigned int min_hits = 2;
  tracker_seeded.setRejectGhosts(true);
  tracker_seeded.setChi2Cut(4.0);
  tracker_seeded.setPrintTimings(true);
  tracker_seeded.setVerbosity(1);
  tracker_seeded.setCutOnDca(false);
  tracker_seeded.setSmoothBack(true);
  
  
  TFile* mcfile = new TFile(argv[2], "recreate");
  TTree* mc_tree = new TTree("mc_events", "a tree of SimpleMCEvent");
  SimpleMCEvent* mcevent = new SimpleMCEvent();
  mc_tree->Branch("tracks", "SimpleMCEvent", &mcevent);
  TTree* reco_tree = new TTree("reco_events", "a tree of SimpleRecoEvent");
  SimpleRecoEvent* recoevent = new SimpleRecoEvent();
  reco_tree->Branch("tracks", "SimpleRecoEvent", &recoevent);
  
  for(unsigned int ev=0;ev<etree->GetEntries();ev++)
  {
    mcevent->tracks.clear();
    recoevent->tracks.clear();
    
    infile.cd();
    vector<double> pt_vec;
    
    etree->GetEntry(ev);
    cout<<"event "<<ev<<":"<<endl<<endl;
    
    vector<SimpleHit3D> hits_0_3;
    vector<SimpleHit3D> hits_4_5;
    vector<SimpleHit3D> hits_all;
    vector<SimpleTrack3D> tracks_seed; // list of space points _and_ helix parameters
    vector<SimpleTrack3D> tracks; // list of space points _and_ helix parameters
    
    vector<SimpleMCTrack>& mctracks = inputevent->tracks;
    
    cout<<"total MC tracks = "<<mctracks.size()<<endl;
    cout<<endl;
    
    // the last track contains noise hits
    for(unsigned int trk=0;trk<(mctracks.size() - 1);trk++)
    {
      pt_vec.push_back(0.003*2./mctracks[trk].kappa);
      if(print_truth==true)
      {
        cout<<"truth track : "<<trk<<endl;
        cout<<"phi = "<<mctracks[trk].phi<<endl;
        cout<<"d = "<<mctracks[trk].d<<endl;
        cout<<"kappa = "<<mctracks[trk].kappa<<endl;
        cout<<"dzdl = "<<mctracks[trk].dzdl<<endl;
        cout<<"z0 = "<<mctracks[trk].z0<<endl;
        cout<<endl;
      }
      mcevent->tracks.push_back(SimpleMCTrack());
      mcevent->tracks.back().hits.clear();
      mcevent->tracks.back().kappa = mctracks[trk].kappa;
      mcevent->tracks.back().dzdl = mctracks[trk].dzdl;
      mcevent->tracks.back().d = mctracks[trk].d;
      mcevent->tracks.back().phi = mctracks[trk].phi;
      mcevent->tracks.back().z0 = mctracks[trk].z0;
      
      vector<SimpleMCHit>& mchits = mctracks[trk].hits;
      
      for(unsigned int hit=0;hit<mchits.size();hit++)
      {
        SimpleMCHit& mchit = mchits[hit];
        float phi = atan2(mchit.y, mchit.x);
        float xy_error = smear_xy_layer[mchit.layer]*sqrt_12*0.5;
        float x_error = fabs(xy_error*sin(phi));
        float y_error = fabs(xy_error*cos(phi));
        float z_error = smear_z_layer[mchit.layer]*sqrt_12*0.5;
        
        mcevent->tracks.back().hits.push_back(SimpleMCHit(mchit.x, mchit.y, mchit.z, mchit.layer, mchit.index));
        
        if(mchit.layer<4)
        {
          hits_0_3.push_back(SimpleHit3D(mchit.x,x_error, mchit.y,y_error, mchit.z,z_error, mchit.index, mchit.layer));
        }
        else
        {
          hits_4_5.push_back(SimpleHit3D(mchit.x,x_error, mchit.y,y_error, mchit.z,z_error, mchit.index, mchit.layer));
        }
        hits_all.push_back(SimpleHit3D(mchit.x,x_error, mchit.y,y_error, mchit.z,z_error, mchit.index, mchit.layer));
        if(print_truth==true)
        {
          cout<<mchit.x<<" "<<mchit.y<<" "<<mchit.z<<endl;
        }
      }
      
      if(print_truth==true){cout<<endl<<endl;}
    }
    mcfile->cd();
    mc_tree->Fill();
    
    vector<SimpleMCHit>& mchits = mctracks.back().hits;
    for(unsigned int h=0;h<mchits.size();h++)
    {
      SimpleMCHit& mchit = mchits[h];
      float phi = atan2(mchit.y, mchit.x);
      float xy_error = smear_xy_layer[mchit.layer]*sqrt_12*0.5;
      float x_error = fabs(xy_error*sin(phi));
      float y_error = fabs(xy_error*cos(phi));
      float z_error = smear_z_layer[mchit.layer]*sqrt_12*0.5;
      
      if(mchit.layer<4)
      {
        hits_0_3.push_back(SimpleHit3D(mchit.x,x_error, mchit.y,y_error, mchit.z,z_error, mchit.index, mchit.layer));
      }
      else
      {
        hits_4_5.push_back(SimpleHit3D(mchit.x,x_error, mchit.y,y_error, mchit.z,z_error, mchit.index, mchit.layer));
      }
      hits_all.push_back(SimpleHit3D(mchit.x,x_error, mchit.y,y_error, mchit.z,z_error, mchit.index, mchit.layer));
    }
    
    
    timeval t1,t2;
    double time1=0.;
    double time2=0.;
    unsigned int ngood = 0;
    unsigned int n_trk = tracks.size();

    tracks.clear();
    tracks_seed.clear();
    
    gettimeofday(&t1, NULL);
    tracker.findHelices(hits_0_3, min_hits_seed, max_hits_seed, tracks_seed);
    gettimeofday(&t2, NULL);
    time1 = ((double)(t1.tv_sec) + (double)(t1.tv_usec)/1000000.);
    time2 = ((double)(t2.tv_sec) + (double)(t2.tv_usec)/1000000.);
    cout<<"seed tracking time = "<<(time2 - time1)<<endl<<endl;
    tracker_seeded.setSeedStates(tracker.getKalmanStates());
    
    cout<<tracks_seed.size()<<" track seeds found"<<endl<<endl;
    
    gettimeofday(&t1, NULL);
    tracker_seeded.findSeededHelices(tracks_seed, hits_4_5, min_hits, max_hits, tracks);
    gettimeofday(&t2, NULL);
    time1 = ((double)(t1.tv_sec) + (double)(t1.tv_usec)/1000000.);
    time2 = ((double)(t2.tv_sec) + (double)(t2.tv_usec)/1000000.);
    cout<<"primary tracking time = "<<(time2 - time1)<<endl;
    vector<double>& isolation = tracker_seeded.getIsolation();
    
    
    
    set<unsigned int> reco_vec;
    
    ngood = 0;
    n_trk = tracks.size();
    for(unsigned int trk=0;trk<n_trk;++trk)
    {
      unsigned int high = 0;
      unsigned int low = (1<<31);
      unsigned int n_hit = tracks[trk].hits.size();
      for(unsigned int ht=0;ht<n_hit;++ht)
      {
        tracks[trk].hits[ht].index > high ? high = tracks[trk].hits[ht].index : high = high;
        tracks[trk].hits[ht].index < low ? low = tracks[trk].hits[ht].index : low = low;
      }
      if( ((high - low) == 5) && ( (low%6)==0 ) )
      {
        ngood++;
        reco_vec.insert(low/6);
      }
    }
    
    cout<<endl;
    
    
    if(print_reco==true)
    {
      for(unsigned int trk=0;trk<n_trk;++trk)
      {
        cout<<"track "<<trk<<" : "<<endl;
        unsigned int n_hit = tracks[trk].hits.size();
        for(unsigned int ht=0;ht<n_hit;++ht)
        {
          cout<<"hit "<<ht<<" : ";
          cout<<tracks[trk].hits[ht].x<<" ";
          cout<<tracks[trk].hits[ht].y<<" ";
          cout<<tracks[trk].hits[ht].z<<" ";
          cout<<tracks[trk].hits[ht].index<<endl;
        }
        cout<<"phi = "<<tracks[trk].phi<<endl;
        cout<<"d = "<<tracks[trk].d<<endl;
        cout<<"kappa = "<<tracks[trk].kappa<<endl;
        cout<<"dzdl = "<<tracks[trk].dzdl<<endl;
        cout<<"z0 = "<<tracks[trk].z0<<endl;
        cout<<endl;
      }
    }
    
    
    cout<<n_trk<<" tracks found"<<endl;
    cout<<ngood<<" good tracks found"<<endl;
    cout<<endl<<endl;
    
    for(unsigned int trk=0;trk<n_trk;++trk)
    {
      recoevent->tracks.push_back(SimpleRecoTrack());
      recoevent->tracks.back().kappa = tracks[trk].kappa;
      recoevent->tracks.back().phi = tracks[trk].phi;
      recoevent->tracks.back().d = tracks[trk].d;
      recoevent->tracks.back().z0 = tracks[trk].z0;
      recoevent->tracks.back().dzdl = tracks[trk].dzdl;
      for(unsigned int h=0;h<tracks[trk].hits.size();++h)
      {
        recoevent->tracks.back().indexes.push_back(tracks[trk].hits[h].index);
      }
      recoevent->tracks.back().chi2 = (tracker_seeded.getKalmanStates())[trk].chi2/7.;
      recoevent->tracks.back().isolation = isolation[trk];
    }
    mcfile->cd();
    reco_tree->Fill();
    
  }
  
  
  mcfile->cd();
  mc_tree->Write();
  reco_tree->Write();
  mcfile->Close();
  mcfile->Delete();
  
  return 0;
}


